﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace DokanNetMirror
{
    class PseudoRC4 : IFileSystemEncryption
    {
        const int STREAMSIZE = 64 * 1024 * 1024;
        const int STREAMMASK = STREAMSIZE - 1;
        byte[] state = new byte[256];
        uint stateIndex1, stateIndex2;
        byte[] randomstream = new byte[STREAMSIZE];
        string swappableString;

        void swap(uint i, uint j)
        {
            byte temp = state[i];
            state[i] = state[j];
            state[j] = temp;
        }

        /* KSA */
        void RC4Init(byte[] key)
        {
            int keyLength = key.Length;
            for (stateIndex1 = 0; stateIndex1 < 256; stateIndex1++)
            {
                state[stateIndex1] = (byte)stateIndex1;
            }

            for (stateIndex1 = stateIndex2 = 0; stateIndex1 < 256; stateIndex1++)
            {
                stateIndex2 = (stateIndex2 + key[stateIndex1 % keyLength] + state[stateIndex1]) & 255;
                swap(stateIndex1, stateIndex2);
            }

            stateIndex1 = stateIndex2 = 0;
        }

        void RC4Fill()
        {
            for (int i = 0; i < STREAMSIZE; i++)
            {
                stateIndex1 = (stateIndex1 + 1) & 255;
                stateIndex2 = (stateIndex2 + state[stateIndex1]) & 255;

                swap(stateIndex1, stateIndex2);
                randomstream[i] = state[(state[stateIndex1] + state[stateIndex2]) & 255];
            }
        }


        #region IFileSystemEncryption Members

        public void Init(object seed)
        {
            string pwd = seed as string ?? "password";
            byte[] pwdBytes = Encoding.Unicode.GetBytes(pwd);
            RC4Init(pwdBytes);
            RC4Fill();
            swappableString = Randomize("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", pwdBytes);
        }

        private static string Randomize(string input, byte[] seed)
        {
            StringBuilder randomchars = new StringBuilder(input);
            int keyLength = seed.Length;
            for (int i = 0; i < randomchars.Length - 1; i++)
            {
                char c = randomchars[i];
                int j = seed[i % keyLength] % (randomchars.Length - i - 1) + i + 1;
                randomchars[i] = randomchars[j];
                randomchars[j] = c;
            }

            string temp = randomchars.ToString();
            for (int i = 0; i < randomchars.Length; i += 2)
            {
                int index1 = input.IndexOf(temp[i]);
                int index2 = input.IndexOf(temp[i + 1]);
                randomchars[index1] = input[index2];
                randomchars[index2] = input[index1];
            }

            return randomchars.ToString();
        }



        public void EncryptBytes(byte[] buffer, long initialOffset)
        {
            long offset = initialOffset;
            for (int i = 0; i < buffer.Length; i++)
            {
                offset = offset & STREAMMASK;
                buffer[i] = (byte)(buffer[i] ^ randomstream[offset]);
                offset++;
            }
        }

        public string EncryptFileName(string fileName)
        {
            StringBuilder result = new StringBuilder(fileName.Length);
            foreach (var c in fileName)
            {
                int index = 0;
                if (c >= 'a' && c <= 'z')
                {
                    index = c - 'a';
                    result.Append(swappableString[index]);
                    continue;
                }
                if (c >= 'A' && c <= 'Z')
                {
                    index = c - 'A' + 26;
                    result.Append(swappableString[index]);
                    continue;
                }
                if (c >= '0' && c <= '9')
                {
                    index = c - '0' + 52;
                    result.Append(swappableString[index]);
                    continue;
                }

                result.Append(c);
            }

            return result.ToString();
        }

        #endregion
    }
}
